# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import experimental.prophet as prophet
from experimental.prophet.craps.universe import universe_all


def signed_log(x):
  r1 = prophet.log(x + 1.)
  r2 = -prophet.log(-x + 1.)
  return prophet.cond(x >= 0, r1, r2)


def single_product_features(symbol_info, feed, price_multiplier=None):
  price_multiplier = price_multiplier or symbol_info.price_multiplier
  windows = ['1s', '2s', '5s', '10s', '30s', '1m']
  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs

  timer_1s = prophet.timer('1s')
  ts = prophet.max(prophet.max(book.timestamp, trade.timestamp), timer_1s)
  with prophet.as_default_timestamp(ts):
    ask0_p = (book.estimated_ask0_price * price_multiplier).named_ref('ask0_p.x')
    bid0_p = (book.estimated_bid0_price * price_multiplier).named_ref('bid0_p.x')
    mid_p = ((ask0_p + bid0_p) / 2.).named_ref('mid_p.x')

    spread = (ask0_p - bid0_p).named_ref('spread')
    spread_mas = prophet.time_moving_window(spread, ['1s', '2s', '5s', '10s']).mean()

    mid_p_stds = []
    for w in windows:
      mid_p_std = prophet.time_moving_window(mid_p, w).std()
      mid_p_std = prophet.cond(prophet.isnan(mid_p_std), 0., mid_p_std).named_ref('mid_p.mstd' + w)
      mid_p_stds.append(mid_p_std)

    mid_rets = []
    mid_ret_logs = []
    for w in windows:
      mid_ret = (mid_p - prophet.time_shift(mid_p, w)).named_ref('mid_ret.' + w)
      mid_rets.append(mid_ret)
      mid_ret_logs.append(signed_log(mid_ret).named_ref('mid_ret.%s.log' % w))

    tp = (trade.price * price_multiplier).named_ref('trade_price')
    if not symbol_info.inverse:
      tq = trade.qty.named_ref('trade_qty')
    else:
      tq = (trade.qty / trade.price).named_ref('trade_qty')
    tside = trade.side.named_ref('trade_side')

    tp_stds = []
    for w in windows:
      tp_std = prophet.time_moving_window(tp, w).std()
      tp_std = prophet.cond(prophet.isnan(tp_std), 0., tp_std).named_ref('trade_price.mstd' + w)
      tp_stds.append(tp_std)

    tp_mdiffs = []
    for w in windows:
      tp_mdiff = prophet.time_moving_window(tp, w).max_min_diff()
      tp_mdiffs.append(tp_mdiff)

    tq_sum_vars = []
    for w in windows[:4]:
      tq_sum = prophet.time_moving_window(tq, w).sum().named_ref('trade_qty.' + w)
      tq_sum_log = signed_log(tq_sum).named_ref('trade_qty.%s.log' % w)
      tq_sum_vars += [tq_sum, tq_sum_log]

    tq_vars = []
    with prophet.control_if(tside == 1):
      for w in windows[:4]:
        buy_tq = prophet.time_moving_window(tq, w).sum().named_ref('buy_trade_qty.' + w)
        buy_tq_log = signed_log(buy_tq).named_ref('buy_trade_qty.%s.log' % w)
        tq_vars += [buy_tq, buy_tq_log]

    with prophet.control_if(tside == 2):
      for w in windows[:4]:
        sell_tq = prophet.time_moving_window(tq, w).sum().named_ref('sell_trade_qty.' + w)
        sell_tq_log = signed_log(sell_tq).named_ref('sell_trade_qty.%s.log' % w)
        tq_vars += [sell_tq, sell_tq_log]

    prev_ask0_p = prophet.shift(ask0_p, 1)
    prev_bid0_p = prophet.shift(bid0_p, 1)
    ask_flip = (ask0_p > prev_ask0_p).named_ref('ask_flip')
    bid_flip = (bid0_p < prev_bid0_p).named_ref('bid_flip')

    return [
        ask0_p,
        bid0_p,
        mid_p,
        spread,
        *spread_mas,
        *mid_p_stds,
        *tp_stds,
        *tp_mdiffs,
        *mid_rets,
        *mid_ret_logs,
        *tq_vars,
        *tq_sum_vars,
        ask_flip,
        bid_flip
    ]


def features(target_coin, additional_symbols, date):
  if target_coin == 'btc':
    symbols = sorted(
        set([
            'bitmex-xbtusd',
            'huobi-btcusd-quarter',
            'okex-btcusd-quarter',
            'okex-btcusd-this_week',
            'okex-btcusd-next_week',
            'okex-btcusd-perpetual',
            'binance-btcusdt',
            'gdax-btcusd',
        ] + additional_symbols))

  elif target_coin == 'eth':
    symbols = sorted(
        set([
            'bitmex-ethusd',
            'huobi-ethusd-quarter',
            'okex-ethusd-quarter',
            'okex-ethusd-this_week',
            'okex-ethusd-next_week',
            'okex-ethusd-perpetual',
            'binance-ethusdt',
            'gdax-ethusd',
        ] + additional_symbols))
    # TODO: add btc

  else:
    raise ValueError(target_coin)

  universe = [s for s in universe_all(date) if s.nickname in symbols]
  vars = []

  for symbol_info in universe:
    feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)
    with prophet.namespace(symbol_info.nickname):
      vars += single_product_features(symbol_info, feed)
  return vars


def samplers(target_coin, feature_prefix):
  graph = prophet.get_default_graph()
  # vol = graph.get_named_variable('%sokex-btcusd-quarter/trade_price.mstd10s' % feature_prefix)
  # low_vol = (vol < 1.5).named_ref('low_vol')  # 60%
  # mid_vol = (1.5 <= vol and vol < 5.).named_ref('mid_vol') # 60-95%
  # high_vol = (5. <= vol).named_ref('high_vol')  # 95%

  if target_coin == 'btc':
    vol = graph.get_named_variable('%sokex-btcusd-quarter/trade_price.mmaxdiff30s' % feature_prefix)
    low_vol = (vol < 8.).named_ref('low_vol')  # 60%
    mid_vol = (8. <= vol and vol < 25.).named_ref('mid_vol')  # 60-95%
    high_vol = (25. <= vol).named_ref('high_vol')  # 95%

  elif target_coin == 'eth':
    vol = graph.get_named_variable('%sokex-ethusd-quarter/trade_price.mmaxdiff30s' % feature_prefix)
    low_vol = (vol < 0.15).named_ref('low_vol')  # 60%
    mid_vol = (0.15 <= vol and vol < 0.45).named_ref('mid_vol')  # 60-95%
    high_vol = (0.45 <= vol).named_ref('high_vol')  # 95%

  else:
    raise ValueError(target_coin)

  return [low_vol, mid_vol, high_vol]


def trigger(target_coin, trigger_symbol_nick, feature_prefix):
  graph = prophet.get_default_graph()

  if target_coin == 'btc':
    trigger_nicks = set([trigger_symbol_nick, 'bitmex-xbtusd', 'okex-btcusd-quarter'])
  elif target_coin == 'eth':
    trigger_nicks = set([trigger_symbol_nick, 'bitmex-ethusd', 'okex-ethusd-quarter'])
  else:
    raise ValueError(target_coin)

  trigger = None
  for nickname in trigger_nicks:
    ask_flip = graph.get_named_variable('%s%s/ask_flip' % (feature_prefix, nickname))
    bid_flip = graph.get_named_variable('%s%s/bid_flip' % (feature_prefix, nickname))
    if trigger is None:
      trigger = ask_flip or bid_flip
    else:
      trigger = trigger or ask_flip or bid_flip

  return trigger.named_ref('trigger')
